package mcps

import (
	"context"
	"fmt"
	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
	"github.com/shirou/gopsutil/process"
	"sort"
	"strconv"
	"strings"
	"time"
)

func ProcessMCPServer() {
	s := server.NewMCPServer(
		"进程管理工具",
		"1.0.0",
	)

	// 添加进程列表工具
	processListTool := mcp.NewTool("process_list",
		mcp.WithDescription("获取系统进程列表"),
		mcp.WithString("sort_by",
			mcp.Description("排序方式: cpu, memory, pid, name"),
		),
		mcp.WithString("filter",
			mcp.Description("进程名称过滤器，支持部分匹配"),
		),
	)

	// 添加进程详情工具
	processInfoTool := mcp.NewTool("process_info",
		mcp.WithDescription("获取指定进程的详细信息"),
		mcp.WithString("pid",
			mcp.Required(),
			mcp.Description("进程ID"),
		),
	)

	// 添加进程资源监控工具
	processMonitorTool := mcp.NewTool("process_monitor",
		mcp.WithDescription("监控指定进程的资源使用情况"),
		mcp.WithString("pid",
			mcp.Required(),
			mcp.Description("进程ID"),
		),
		mcp.WithString("duration",
			mcp.Description("监控持续时间（如：30s, 1m, 5m），默认为1分钟"),
		),
	)

	// 添加工具处理函数
	s.AddTool(processListTool, processListHandler)
	s.AddTool(processInfoTool, processInfoHandler)
	s.AddTool(processMonitorTool, processMonitorHandler)

	// 启动服务器
	fmt.Println("进程管理服务已启动...")
	if err := server.ServeStdio(s); err != nil {
		fmt.Printf("服务器错误: %v\n", err)
	}
}

func processListHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	// 获取所有进程
	processes, err := process.Processes()
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("获取进程列表失败: %v", err)), nil
	}

	// 获取过滤器
	filter, _ := request.Params.Arguments["filter"].(string)

	// 收集进程信息
	var processInfos []map[string]interface{}
	for _, p := range processes {
		name, err := p.Name()
		if err != nil {
			continue
		}

		// 应用过滤器
		if filter != "" && !strings.Contains(strings.ToLower(name), strings.ToLower(filter)) {
			continue
		}

		cpu, _ := p.CPUPercent()
		mem, _ := p.MemoryPercent()

		processInfos = append(processInfos, map[string]interface{}{
			"pid":      p.Pid,
			"name":     name,
			"cpu_used": cpu,
			"mem_used": mem,
		})
	}

	// 排序
	sortBy, _ := request.Params.Arguments["sort_by"].(string)
	switch sortBy {
	case "cpu":
		sort.Slice(processInfos, func(i, j int) bool {
			return processInfos[i]["cpu_used"].(float64) > processInfos[j]["cpu_used"].(float64)
		})
	case "memory":
		sort.Slice(processInfos, func(i, j int) bool {
			return processInfos[i]["mem_used"].(float32) > processInfos[j]["mem_used"].(float32)
		})
	case "name":
		sort.Slice(processInfos, func(i, j int) bool {
			return processInfos[i]["name"].(string) < processInfos[j]["name"].(string)
		})
	case "pid":
		sort.Slice(processInfos, func(i, j int) bool {
			return processInfos[i]["pid"].(int32) < processInfos[j]["pid"].(int32)
		})
	}

	// 格式化输出
	var result strings.Builder
	result.WriteString("===== 进程列表 =====\n")
	result.WriteString(fmt.Sprintf("总进程数: %d\n\n", len(processInfos)))
	result.WriteString("PID\t进程名\tCPU使用率\t内存使用率\n")
	result.WriteString("----------------------------------------\n")

	for _, info := range processInfos {
		result.WriteString(fmt.Sprintf("%d\t%s\t%.2f%%\t%.2f%%\n",
			info["pid"].(int32),
			info["name"].(string),
			info["cpu_used"].(float64),
			info["mem_used"].(float32),
		))
	}

	return mcp.NewToolResultText(result.String()), nil
}

func processInfoHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	// 获取进程ID
	pidStr, ok := request.Params.Arguments["pid"].(string)
	if !ok {
		return mcp.NewToolResultError("pid 参数必须是字符串类型"), nil
	}

	pid, err := strconv.Atoi(pidStr)
	if err != nil {
		return mcp.NewToolResultError("无效的进程ID"), nil
	}

	// 获取进程信息
	p, err := process.NewProcess(int32(pid))
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("进程不存在: %v", err)), nil
	}

	// 收集进程详细信息
	name, _ := p.Name()
	cmdline, _ := p.Cmdline()
	cwd, _ := p.Cwd()
	createTime, _ := p.CreateTime()
	cpu, _ := p.CPUPercent()
	mem, _ := p.MemoryPercent()
	memInfo, _ := p.MemoryInfo()
	status, _ := p.Status()

	// 格式化输出
	var result strings.Builder
	result.WriteString(fmt.Sprintf("===== 进程详情 (PID: %d) =====\n", pid))
	result.WriteString(fmt.Sprintf("进程名称: %s\n", name))
	result.WriteString(fmt.Sprintf("命令行: %s\n", cmdline))
	result.WriteString(fmt.Sprintf("工作目录: %s\n", cwd))
	result.WriteString(fmt.Sprintf("创建时间: %s\n", time.Unix(createTime/1000, 0).Format("2006-01-02 15:04:05")))
	result.WriteString(fmt.Sprintf("进程状态: %s\n", status))
	result.WriteString(fmt.Sprintf("CPU使用率: %.2f%%\n", cpu))
	result.WriteString(fmt.Sprintf("内存使用率: %.2f%%\n", mem))
	if memInfo != nil {
		result.WriteString(fmt.Sprintf("内存使用详情:\n"))
		result.WriteString(fmt.Sprintf("  - RSS: %.2f MB\n", float64(memInfo.RSS)/(1024*1024)))
		result.WriteString(fmt.Sprintf("  - VMS: %.2f MB\n", float64(memInfo.VMS)/(1024*1024)))
	}

	return mcp.NewToolResultText(result.String()), nil
}

func processMonitorHandler(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	// 获取进程ID
	pidStr, ok := request.Params.Arguments["pid"].(string)
	if !ok {
		return mcp.NewToolResultError("pid 参数必须是字符串类型"), nil
	}

	pid, err := strconv.Atoi(pidStr)
	if err != nil {
		return mcp.NewToolResultError("无效的进程ID"), nil
	}

	// 获取监控持续时间
	duration := "1m"
	if d, ok := request.Params.Arguments["duration"].(string); ok && d != "" {
		duration = d
	}

	// 解析持续时间
	d, err := time.ParseDuration(duration)
	if err != nil {
		return mcp.NewToolResultError("无效的持续时间格式"), nil
	}

	// 获取进程信息
	p, err := process.NewProcess(int32(pid))
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("进程不存在: %v", err)), nil
	}

	// 开始监控
	name, _ := p.Name()
	var result strings.Builder
	result.WriteString(fmt.Sprintf("===== 进程监控 (PID: %d, 名称: %s) =====\n", pid, name))
	result.WriteString(fmt.Sprintf("监控时长: %s\n\n", duration))
	result.WriteString("时间\t\tCPU使用率\t内存使用率\n")
	result.WriteString("----------------------------------------\n")

	start := time.Now()
	for time.Since(start) < d {
		cpu, _ := p.CPUPercent()
		mem, _ := p.MemoryPercent()

		result.WriteString(fmt.Sprintf("%s\t%.2f%%\t\t%.2f%%\n",
			time.Now().Format("15:04:05"),
			cpu,
			mem,
		))

		time.Sleep(time.Second)
	}

	return mcp.NewToolResultText(result.String()), nil
}
